import numpy as np
import torch
from scipy import stats
import cv2
import os
import argparse
from torchvision import transforms
import sys

# 导入自定义模型
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.vgg import VGGNet

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='评估 TID2013 数据集上的 RankIQA 模型')
    parser.add_argument('--model_path', type=str, 
                        default='./checkpoints/rank_tid2013/model_best.pth.tar',
                        help='预训练模型路径')
    parser.add_argument('--test_file', type=str, 
                        default='../data/score_tid2013/tid2013_test.txt',
                        help='测试文件路径')
    parser.add_argument('--res_dir', type=str, 
                        default='./results/tid2013/',
                        help='结果保存目录')
    parser.add_argument('--num_patch', type=int, 
                        default=30,
                        help='每张图像的随机裁剪数量')
    return parser.parse_args()

def load_model(model_path):
    """加载预训练的模型"""
    # 创建 VGG 模型
    model = VGGNet(pretrained=False)
    
    # 加载预训练权重
    checkpoint = torch.load(model_path)
    if 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['state_dict'])
    else:
        model.load_state_dict(checkpoint)
    
    model = model.to(device)
    model.eval()
    return model

def main():
    """主函数"""
    args = parse_args()
    
    # 确保结果目录存在
    os.makedirs(args.res_dir, exist_ok=True)
    
    # 加载模型
    model = load_model(args.model_path)
    
    # 打开结果文件
    tp = 'Rank_all'
    srocc_file = open(os.path.join(args.res_dir, tp + '_srocc.txt'), "w")
    lcc_file = open(os.path.join(args.res_dir, tp + '_lcc.txt'), "w")
    
    # 读取测试文件
    filename = [line.rstrip('\n') for line in open(args.test_file)]
    
    roidb = []
    scores = []
    for i in filename:
        roidb.append(i.split()[0])
        scores.append(float(i.split()[1]))
    scores = np.asarray(scores)
    
    # 图像预处理
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                             std=[0.229, 0.224, 0.225])
    ])
    
    # 评估
    Num_Patch = args.num_patch
    Num_Image = len(scores)
    feat = np.zeros([Num_Image, Num_Patch])
    pre = np.zeros(Num_Image)
    med = np.zeros(Num_Image)
    
    with torch.no_grad():
        for i in range(Num_Image):
            directory = roidb[i]
            im = cv2.imread(directory)
            if im is None:
                print(f"无法读取图像: {directory}")
                continue
                
            # 转换为 RGB
            im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
            
            for j in range(Num_Patch):
                x = im.shape[0]
                y = im.shape[1]
                x_p = np.random.randint(x-224, size=1)[0]
                y_p = np.random.randint(y-224, size=1)[0]
                
                # 裁剪并转换图像
                patch = im[x_p:x_p+224, y_p:y_p+224, :]
                patch_tensor = transform(patch).unsqueeze(0).to(device)
                
                # 前向传播
                output = model(patch_tensor)
                
                # 保存结果
                feat[i, j] = output.item()
                pre[i] += output.item()
            
            pre[i] /= Num_Patch
            med[i] = np.median(feat[i, :])
    
    # 计算相关系数
    srocc = stats.spearmanr(pre, scores)[0]
    lcc = stats.pearsonr(pre, scores)[0]
    print(f'LCC of mean: {lcc:.4f}')
    print(f'SROCC of mean: {srocc:.4f}')
    
    srocc_file.write('%6.3f\n' % srocc)
    lcc_file.write('%6.3f\n' % lcc)
    srocc_file.close()
    lcc_file.close()
    
    # 计算中位数结果的相关系数
    srocc = stats.spearmanr(med, scores)[0]
    lcc = stats.pearsonr(med, scores)[0]
    print(f'LCC of median: {lcc:.4f}')
    print(f'SROCC of median: {srocc:.4f}')

if __name__ == '__main__':
    main()